import os
import time
import cv2
import imageio
import numpy as np


# 边缘提取
def getEdges(img):
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 图像转换为灰度图
    blur_gray = cv2.GaussianBlur(gray, (5, 5), 0, 0)  # 使用高斯模糊去噪
    edges = cv2.Canny(blur_gray, 30, 60)  # 使用Canny进行边缘检测
    return edges

def judge(img):
    answer = False
    img[:int(0.3*img.shape[0])] = 255
    img[int(0.85 * img.shape[0]):] = 255
    img[img > 255] = 255
    img_back = cv2.bitwise_not(img.astype('uint8'))

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    img_erode = cv2.erode(img_back, kernel)


    cv2.imshow('img2', img_erode)
    ret, labels, stats, centroid = cv2.connectedComponentsWithStats(img_erode)
    max_area = sorted(stats, key = lambda s:s[-1], reverse=False)[-2]
    for i, stat in enumerate(stats):
        cv2.rectangle(img, (stat[0], stat[1]), (stat[0] + stat[2], stat[1] + stat[3]), (25, 25, 255), 3)
    cv2.imshow("result", img_erode)  # 显示图片
    cv2.waitKey(0)
    # 搜索图像中的连通区域
    ret, labels, stats, centroid = cv2.connectedComponentsWithStats(img)
    # for i, stat in enumerate(stats):
    #     if stat[4] > 1000:
    #         # 绘制连通区域
    #         cv2.rectangle(img, (stat[0], stat[1]), (stat[0] + stat[2], stat[1] + stat[3]), (25, 25, 255), 10)
    #         # 按照连通区域的索引来打上标签
    #         # cv2.putText(img, str(stat[4]), (stat[0], stat[1] + 25), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 25, 25), 2)
    #     cv2.imshow('1', img)
    #     cv2.waitKey(1)
    # if 'wriper' in img:
    #     answer = True
    #     return answer


new = False
COUNT_MAX = 1475
count_now = 0
all_edge = np.zeros((576, 704), dtype=float)
frame_list = []
# 执行函数
def do(img):
    # img = cv2.resize(img, (352, 288))
    hight, width, _ = img.shape
    # img = cv2.medianBlur(img[25:258,10:342], 3)
    global new, COUNT_MAX, count_now, all_edge, frame_list

    if new:
        imageio.mimsave('new/result'+str(int(time.time()))+'.gif', frame_list, 'GIF', duration=0.01)
        all_edge = np.zeros((img.shape[0], img.shape[1]), dtype=float)
        frame_list = []
        new = False

    edge = getEdges(img)
    all_edge += edge
    all_edge[all_edge > 0] = 255
    all_edge = all_edge.astype('uint8')
    each_frame = cv2.merge([all_edge, all_edge, all_edge])
    frame_list.append(each_frame)
    # cv2.imshow('edge', edge)
    cv2.imshow('all', frame_list[-1])
    cv2.waitKey(1)

    count_now += 1
    print(count_now)
    if (new and count_now>0) or count_now >= COUNT_MAX :
        imageio.mimsave('new/result'+str(int(time.time()))+'.gif', frame_list, 'GIF', duration=0.000001)
        all_edge[all_edge > 255] = 255
        all_edge_fanzhuan = 255 - all_edge
        # all_edge_fanzhuan = cv2.bitwise_not(img.astype('uint8'))

        contours, hierarchy = cv2.findContours(all_edge_fanzhuan.astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            ares = int(cv2.contourArea(contour))  # 计算包围形状的面积
            if ares < 1500:  # 过滤面积小于ares的形状
                continue
            """传入一个轮廓图像，返回 x y 是左上角的点， w和h是矩形边框的宽度和高度"""
            x, y, w, h = cv2.boundingRect(contour)
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 3)
            """画出矩形 img 是要画出轮廓的原图(x, y) 是左上角点的坐标(x+w, y+h) 是右下角的坐标（0,255,0）是画线对应的rgb颜色 2 是画出线的宽度 """
            # 获得最小的矩形轮廓 可能带旋转角度
            rect = cv2.minAreaRect(contour)
            jd = abs(rect[2] - 90)  # 旋转角度
            hd = int(rect[1][0])  # 矩形轮廓高度
            wd = int(rect[1][1])  # 矩形轮廓宽度
            zx_x = int(rect[0][0])
            zx_y = int(rect[0][1])
            ab = int((ares * 100) / (wd * hd))  # 遮挡物所占百分比
            if 51 < ab < 78 and 40 < wd < 90 and ((50 < jd < 85 and hd > 2.8 * wd and x < 25 and w > 0.45 * width and (
                    x + w > 0.98 * width or y + h > 0.98 * hight)) or (
                                                          5 < jd < 50 and hd > 2.4 * wd and y < 25 and h > 0.95 * hight and x + w > 0.7 * width)):
                output_name = '123.jpg'

                detectresultpath = 'new/'

                a = [output_name, ab, ares, [x, y, w, h], rect]
                with open('./occlusion1.txt', 'a+') as f:
                    f.write(str(a) + ' \n')
                # 计算最小区域的坐标
                box = cv2.boxPoints(rect)
                # 坐标规范化为整数
                box = np.int0(box)
                cv2.drawContours(img, contour, -1, (0, 255, 255), 2)  # 绘制边缘轮廓
                output = cv2.drawContours(img, [box], 0, (0, 0, 255), 2)  # 画出轮廓
                cv2.imwrite(os.path.join(detectresultpath, '{}'.format(output_name)), all_edge_fanzhuan)
                return True
            elif wd < 120 and ((50 < jd < 85 and hd > 2.4 * wd and x < 25 and w > 0.6 * width and (
                    x + w > 0.98 * width or y + h > 0.98 * hight)) or (
                                       5 < jd < 50 and hd > 2.4 * wd and y < 25 and h > 0.95 * hight and x + w > 0.7 * width)):
                output_name = '123.jpg'
                detectresultpath = 'new/'
                a = [output_name, ab, ares, [x, y, w, h], rect]
                with open('./occlusion2.txt', 'a+') as f:
                    f.write(str(a) + ' \n')
                # 计算最小区域的坐标
                box = cv2.boxPoints(rect)
                # 坐标规范化为整数
                box = np.int0(box)
                cv2.drawContours(img, contour, -1, (0, 255, 255), 2)  # 绘制边缘轮廓
                output = cv2.drawContours(img, [box], 0, (0, 0, 255), 2)  # 画出轮廓
                cv2.imwrite(os.path.join(detectresultpath, '{}'.format(output_name)), all_edge_fanzhuan)
                return True
            else:
                output_name = '123.jpg'
                detectresultpath = 'new/'
                a = [output_name, ab, ares, [x, y, w, h], rect]
                with open('./occlusion2.txt', 'a+') as f:
                    f.write(str(a) + ' \n')
                # 计算最小区域的坐标
                box = cv2.boxPoints(rect)
                # 坐标规范化为整数
                box = np.int0(box)
                cv2.drawContours(img, contour, -1, (0, 255, 255), 2)  # 绘制边缘轮廓
                output = cv2.drawContours(img, [box], 0, (0, 0, 255), 2)  # 画出轮廓
                cv2.imwrite(os.path.join(detectresultpath, '{}'.format(output_name)), output)
                print()
                count_now = 0
                return

        new = True
        count_now = 0
        # return judge(all_edge)



path = r'C:\Users\29626\Desktop\Rail-detection-master/'
for video in os.listdir(path):
    new = True
    if '司机室' in video:
        continue
    if not video.endswith('.mp4'):
        continue

    capture = cv2.VideoCapture(path + video)
    # print('帧率', capture.get(cv2.CAP_PROP_FRAME_COUNT))
    frame, count_now = 0, 0
    if capture.isOpened():
        while True:
            ret, img = capture.read()  # img 就是一帧图片
            if not ret:
                break  # 当获取完最后一帧就结束

            # if count == 0:
            #     height, width, _ = img.shape
            #     all_edge = np.zeros((height, width), dtype=float)

            frame += 1
            if frame % 1 == 0:
                do(img)
                # print(all_edge)

                # ret, img = capture.read()  # img 就是一帧图片
                # if not ret:
                #     video = video.replace("北京蓝天多维_", "").replace("一端路况_", "").replace("二端路况_", "")
                #     cv2.imwrite('result/'+video[:-4]+'_'+str(frame)+'.jpg', all_edge)
                #     count=0
                #     break  # 当获取完最后一帧就结束
                # if count >= 900:
                #     video = video.replace("北京蓝天多维_", "").replace("一端路况_", "").replace("二端路况_", "")
                #     cv2.imwrite('result/'+video[:-4]+'_'+str(frame)+'.jpg', all_edge)
                #     count = 0
                #     print('时间', time.time() - t1)
    count = 0
